!%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
!
! The `halofit' code models the nonlinear evolution of cold matter 
! cosmological power spectra. The full details of the way in which 
! this is done are presented in Smith et al. (2002), MNRAS, ?, ?. 
! astro-ph/0207664
!
! The code `halofit' was written by R. E. Smith & J. A. Peacock. 
!
! Last edited 8/5/2002.
! DEREK: Interfaced with my code 8/16/2.
!
!%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

!**************************************************************
! A module to generate a fit of the Gaussian smoothed variance.
! It fits a Chebyshev polynomial to the Gaussian smoothed variance, so that you can easily fit the derivative function by
! calling chebyshevDerivative( getFit(GSV) )
!
! type(tGSV) :: GSV
! new( GSV, [lps] ) --- constructor
!
! evaluate(GSV, R)
!
!**************************************************************
module mGSV

        use mPrecision
        use mObject
        use mLPS
        use mFit

        implicit none

        real(D), parameter :: LL_R = 1e-3_D, &
                              UL_R = 1e4_D

        interface new
                module procedure new_GSV
        end interface new

        interface write
                module procedure write_GSV
        end interface write

        interface getUniverse
                module procedure getUniverse_GSV
        end interface getUniverse

        interface evaluate
                module procedure evaluate_GSV
        end interface evaluate

        interface getLPS
                module procedure getLPS_GSV
        end interface getLPS

        interface getH
                module procedure getH_GSV
        end interface getH

        interface fitDone
                module procedure fitDone_GSV
        end interface fitDone

        private
        public :: new, write
        public :: getUniverse, getFit, evaluate, getLPS, fitDone, getH

contains

        subroutine new_GSV(gsv, lps)
                type(tGSV), intent(out) :: gsv
                type(tLPS), intent(in), optional :: lps

                if( present(lps) ) then
                        gsv%lps = lps
                else
                        call new(gsv%lps)
                end if

                call new(gsv%GSVFit, LL_R, UL_R, xLog = .true., yLog = .true., lps = gsv%lps)

                call fit(gsv%GSVfit, GSVFunction)
        end subroutine new_GSV

        real(D) function GSVFunction(params, R)
                type(tFitParams) :: params
                real(D), intent(in) :: R ! Units: Mpc
                GSVFunction = smoothedVariance( getLPS(params), R, 1.0_D, 'GA' )
        end function GSVFunction

        type(tUniverse) pure function getUniverse_GSV(GSV)
                type(tGSV), intent(in) :: GSV
                getUniverse_GSV = getUniverse(GSV%lps)
        end function getUniverse_GSV

        real(D) function getH_GSV(GSV)
                type(tGSV), intent(in) :: GSV
                getH_GSV = getH( getLPS(GSV) )
        end function getH_GSV

        type(tCFit) function getFit(GSV)
                type(tGSV), intent(in) :: GSV
                getFit = GSV%GSVFit
        end function getFit

        real(D) function evaluate_GSV(GSV, x)
                type(tGSV), intent(in) :: GSV
                real(D), intent(in) :: x
                evaluate_GSV = evalChebyshevFunction(GSV%GSVFit, x)
        end function evaluate_GSV

        type(tLPS) pure function getLPS_GSV(GSV)
                type(tGSV), intent(in) :: GSV
                getLPS_GSV = GSV%lps
        end function getLPS_GSV

        logical pure function fitDone_GSV(GSV)
                type(tGSV), intent(in) :: GSV
                fitDone_GSV = cFitDone(GSV%GSVFit)
        end function fitDone_GSV

        subroutine write_GSV(GSV)
                type(tGSV), intent(in) :: GSV
                call writeCFitCoeffs(GSV%GSVFit)
        end subroutine write_GSV

end module mGSV

!*******************************************************************************************************************

! Module mPSFittingFunctions currently has Smith et al., and Peacock & Dodds non-linear fits to the power spectrum
! from simulations.
!
! To use smith, call the constructor:
!
! call new( smith, [lps], [a] )
!
! and evaluate via:
!
! call smithEtAl(smith, k, a, answer)
!
! For Peacock & Dodds, there is only the lps object:
!
! answer = peacockDodds(lps, k, a)
!
!*******************************************************************************************************************

module mPSFittingFunctions

        use mPrecision
        use mObject
        use mGSV

        implicit none

        real(D), parameter :: ROOT_MIN = -5.0_D, &
                              ROOT_MAX = 3.5_D

        interface new
                module procedure new_smith
        end interface new

        interface getUniverse
                module procedure getUniverse_smith
        end interface getUniverse

        interface getLPS
                module procedure getLPS_smith
        end interface getLPS

        interface getH
                module procedure getH_smith
        end interface getH

        interface fitDone
                module procedure fitDone_smith
        end interface fitDone

        private
        public :: new, getUniverse
        public :: smithEtAl
        public :: smithSmoothed
        public :: peacockDodds

contains

        subroutine new_smith(smith, lps, a)
                use mLPS
                use mFit
                type(tSmith), intent(out) :: smith
                type(tLPS), intent(in), optional :: lps
                real(D), intent(in), optional :: a

                if( present(lps) ) then
                        call new(smith%GSV, lps)
                else
                        call new(smith%GSV)
                end if

                smith%dGSV = chebyshevDerivative( getFit(smith%GSV) )
                smith%ddGSV = chebyshevDerivative(smith%dGSV)

!                smith%a = -1.0_D
                if( present(a) ) then
                        call smithSpectralQuantities(smith, a)
                else
                        call smithSpectralQuantities(smith, 1.0_D)
                end if
        end subroutine new_smith

        type(tUniverse) pure function getUniverse_smith(smith)
                type(tSmith), intent(in) :: smith
                getUniverse_smith = getUniverse(smith%GSV)
        end function getUniverse_smith

        type(tLPS) pure function getLPS_smith(smith)
                type(tSmith), intent(in) :: smith
                getLPS_smith = getLPS(smith%GSV)
        end function getLPS_smith

        real(D) function getH_smith(smith)
                type(tSmith), intent(in) :: smith
                getH_smith = getH(smith%GSV)
        end function getH_smith

        subroutine smithEtAl(smith, k, a, ans)

                use mConstants
                use mUniverse
                use mLPS
                use mRoot

                type(tSmith), intent(inout) :: smith
                real(D), intent(in) :: k, & ! wavenumber in units of 1 / Mpc
                                       a    ! scale factor
                real(D), intent(out) :: ans

                type(tRoot) :: root
                real(D) :: growth, plin, pnl, pq, ph

                if(a /= smith%a) call smithSpectralQuantities(smith, a)

                ! linear power spectrum
                ! Remember => plin = k^3 * P(k) * constant,
                ! where constant = 4*pi*V/(2*pi)^3 
                growth = growthFactor( getUniverse(smith), a )
                plin = k ** 3 / (2.0_D * PI * PI) * growth * growth * P_L( getLPS(smith), k, a)

                if(smith%useLinear) then
                        ans = plin
                        return
                else
                        ! calculate nonlinear power according to halofit: pnl = pq + ph,
                        ! where pq represents the quasi-linear (halo-halo) power and 
                        ! where ph is represents the self-correlation halo term. 
                        call halofit( k / getH(smith), &
                                      smith%rneff, smith%rncur, &
                                      smith%rknl, plin, pnl, pq, ph, &
                                      smith%om_m, smith%om_v)

                        ans = pnl
                end if

        end subroutine smithEtAl

        logical pure function fitDone_smith(smith)
                type(tSmith), intent(in) :: smith
                fitDone_smith = fitDone(smith%GSV)
        end function fitDone_smith

        !********************************************************************
        ! Calculate nonlinear wavenumber (rknl), effective spectral
        ! index (rneff), and curvature (rncur) of the power spectrum at
        ! the desired redshift, using method described in Smith et al (2002):
        !********************************************************************
        subroutine smithSpectralQuantities(smith, a)

                use mUniverse
                use mRoot

                type(tSmith), intent(inout) :: smith
                real(D), intent(in) :: a

                type(tRoot) :: root

                call new(root, ROOT_MAX, ROOT_MIN, params = (/ a /), smith = smith)

                if( .not.findRoot(root, logApproxGaussSmoothedVariance, smith%rmid) ) then

                        ! If there is not root, it seems we have a cosmology/redshift that never goes non-linear.
                        ! In this case, have smith just return the linear power.
                        smith%useLinear = .true.
                else

                        smith%useLinear = .false.

                        smith%rmid = exp(smith%rmid)

                        smith%rneff = -3.0_D - log( evaluate( smith%dGSV, smith%rmid) )

                        smith%rncur = - log( evaluate( smith%ddGSV, smith%rmid) )

                        smith%rknl = 1.0_D / ( smith%rmid * getH(smith) )

                        call omegas( getUniverse(smith), a, smith%om_m, smith%om_v)

                end if

                smith%a = a

        end subroutine smithSpectralQuantities

        real(D) function SmithSmoothed(smith, R, a)

                use mIntegration

                type(tSmith), intent(in) :: smith

                real(D), parameter :: LL = 1e-5_D, & ! Lower cutoff
                                      UL = 1e4_D     ! Upper cutoff

                real(D), parameter :: EPS = epsilon(1.0_R4)

                real(D) :: R, a
                type(tIntegral) :: svIntegral

                call new(svIntegral, LL, UL, EPS, changeVariables='EXP', params = (/ R, a /), smith = smith )
                call integrate(svIntegral, svIntegrand, smithSmoothed)

                call checkConvergence(svIntegral, 'SmoothedVariance integration failed to converge.')

        end function smithSmoothed

        !*****************************************************
        ! The integrand for calculating the smoothed variance.
        !*****************************************************
        real(D) function SVIntegrand(integral, k)

                use mConstants
                use mLPS
                use mIntegration

                type(tIntegral), intent(in) :: integral

                real(D), intent(in) :: k ! wavenumber
                real(D) :: R, a, kR, window
                type(tSmith) :: tempSmith

                call getParams(integral, R, a)

                kR = k * R
                window = topHatWindow(kR)
                window = k * window

                tempSmith = getSmith(integral)

                call smithEtAl(tempSmith, k, a, SVIntegrand)
                SVIntegrand = window * window / k ** 3 * SVIntegrand

        end function SVIntegrand

        real(D) function peacockDodds(lps, k, a)

                use mRoot
                use mUniverse

                type(tLPS), intent(in) :: lps
                real(D), intent(in) :: k, a

!                real(D), parameter :: ROOT_MIN = 1e-4_D, &
!                                      ROOT_MAX = 1e4_D
                real(D), parameter :: ROOT_MIN = 1e-10_D, &
                                      ROOT_MAX = 1e4_D

                real(D) :: pnl_pd, growth
                type(tRoot) :: root

                call new( root, ROOT_MAX, ROOT_MIN, params = (/ k, a /), lps = lps )
                if( .not.findRoot(root, rootPeacockDodds, peacockDodds) ) stop

                growth = growthFactor( getUniverse(lps), a )
                peacockDodds = growth * growth * getOutput(root)

        end function peacockDodds

        real(D) function rootPeacockDodds(root, k)

                use mRoot

                type(tRoot), intent(inout) :: root
                real(D), intent(in) :: k

                real(D) :: k_L, a, k_NL, pnl_pd
                type(tLPS) :: lps

                call getParams(root, k_L, a)
                lps = getLPS(root)

                k_NL = k
                call subPeacockDodds(lps, k_NL, a, pnl_pd)
                rootPeacockDodds = k_NL - k_L
                call setOutput(root, pnl_pd)

        end function rootPeacockDodds

        subroutine subPeacockDodds(lps, k, a, pnl_pd)

                use mConstants
               use mLPS

                type(tLPS), intent(in) :: lps
                real(D), intent(inout) :: k ! Input linear wavenumber (1/Mpc)
                                            ! Returns non-linear wavenumber (1/Mpc)
                real(D), intent(in) :: a       ! scale factor
                real(D), intent(out) :: pnl_pd ! Non-linear power
                                               ! a la Peacock & Dodds (1996)
                                               ! k ** 3 / 2pi ** 2 * P_nl(k)
                real(D) :: om_m, om_v
                real(D) :: plin, rn_pd

                plin = k ** 3 / (2.0_D * PI * PI) * P_L(lps, k, a)

                k = k / getH(lps)

                ! effective spectral index: rn_pd=dlogP(k/2)/dlogk          
                rn_pd = rn_cdm(lps, k)

                ! nonlinear power from linear power
                pnl_pd = f_pd(lps, plin, rn_pd, a)

                ! scaling for nonlinear wavenumber
                k = k * getH(lps) * (1.0_D + pnl_pd) ** (1.0_D / 3.0_D)

        end subroutine subPeacockDodds

        !%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        ! halo model nonlinear fitting formula as described in 
        ! Appendix C of Smith et al. (2002)
        subroutine halofit(rk, rn, rncur, rknl, plin, pnl, pq, ph, om_m, om_v)
                real(D), intent(in) :: rk, & ! units: h / Mpc
                                       rn, rncur, &
                                       rknl, & ! units: h/Mpc
                                       plin, om_m, om_v
                real(D), intent(out) :: pnl, pq, ph
                real(D) :: gam, a, b, c, xmu, xnu, alpha, beta, f1, f2, f3
                real(D) :: f1a, f2a, f3a, f1b, f2b, f3b, frac, y
                gam = 0.86485_D + 0.2989_D * rn + 0.1631_D * rncur
                a = 1.4861_D + 1.83693_D * rn + 1.67618_D * rn * rn &
                    + 0.7940_D * rn * rn * rn &
                    + 0.1670756_D * rn * rn * rn * rn - 0.620695_D * rncur
                a = 10 ** a
                b = 10 ** (0.9463_D + 0.9466_D * rn + 0.3084_D * rn * rn &
                           - 0.940_D * rncur)
                c = 10 ** (-0.2807_D + 0.6669_D * rn + 0.3214_D * rn * rn &
                           - 0.0793_D * rncur)
                xmu = 10 ** (-3.54419_D + 0.19086_D * rn)
                xnu = 10 ** (0.95897_D + 1.2857_D * rn)
                alpha = 1.38848_D + 0.3701_D * rn - 0.1452_D * rn * rn
                beta = 0.8291_D + 0.9854_D * rn + 0.3400_D * rn ** 2

                if( abs(1.0_D - om_m) .gt. 0.01_D ) then ! omega evolution 
                        f1a = om_m ** (-0.0732_D)
                        f2a = om_m ** (-0.1423_D)
                        f3a = om_m ** (0.0725_D)
                        f1b = om_m ** (-0.0307_D)
                        f2b = om_m ** (-0.0585_D)
                        f3b = om_m ** (0.0743_D)       
                        frac = om_v / (1.0_D - om_m) 
                        f1 = frac * f1b + (1.0_D - frac) * f1a
                        f2 = frac * f2b + (1.0_D - frac) * f2a
                        f3 = frac * f3b + (1.0_D - frac) * f3a
                else         
                        f1 = 1.0_D
                        f2 = 1.0_D
                        f3 = 1.0_D
                endif

                y = rk / rknl

                ph = a * y ** (f1 * 3.0_D) &
                     / (1.0_D + b * y ** (f2) + (f3 * c * y) ** (3.0_D - gam) )
                ph = ph / ( 1.0_D + xmu * y ** (-1) + xnu * y ** (-2) )
                pq = plin * (1.0_D + plin) ** beta / (1.0_D + plin * alpha) &
                     * exp(-y / 4.0_D - y ** 2 / 8.0_D)

                pnl = pq + ph
        end subroutine halofit

        !%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        ! effective spectral index used in peacock & dodds (1996)
        real(D) function rn_cdm(lps, rk)

                use mConstants
                use mLPS

                type(tLPS), intent(in) :: lps
                real(D), intent(in) :: rk

                real(D) :: yplus, y

                y = rk * getH(lps) / 2.0_D
                y = y ** 3 / (2.0_D * PI * PI) * P_L(lps, y, 1.0_D)
                yplus = rk * getH(lps) * 1.01_D / 2.0_D
                yplus = yplus ** 3 / (2.0_D * PI * PI) * P_L(lps, yplus, 1.0_D)
                rn_cdm = -3.0_D + log(yplus / y) * 100.5_D

        end function rn_cdm

        !%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        ! peacock & dodds (1996) nonlinear fitting formula
        real(D) function f_pd(lps, y, rn, scale_factor)

                use mLPS

                type(tLPS), intent(in) :: lps

                real(D), intent(in) :: y, rn, scale_factor

                real(D) :: om_m, om_v
                real(D) :: g, a, b, alp, bet, vir

                ! calculate matter density, vacuum density at desired redshift
                call omegas(lps, scale_factor, om_m, om_v)

                g = (5.0_D / 2.0_D) * om_m &
                        / ( om_m ** (4.0_D / 7.0_D) - om_v &
                            + (1.0_D + om_m / 2.0_D) * (1 + om_v / 70.0_D) )
                a = 0.482_D * (1.0_D + rn / 3.0_D) ** (-0.947_D)
                b = 0.226_D * (1.0_D + rn / 3.0_D) ** (-1.778_D)
                alp = 3.310_D * (1.0_D + rn / 3.0_D) ** (-0.244_D)
                bet = 0.862_D * (1.0_D + rn / 3.0_D) ** (-0.287_D)
                vir = 11.55_D * (1.0_D + rn / 3.0_D) ** (-0.423_D)
                f_pd = y * ( (1.0_D + b * y * bet + (a * y) ** (alp * bet) ) &
                    / (1.0_D + ( (a * y) ** alp &
                                 * g ** 3 / ( vir * y ** 0.5_D ) ) ** bet) ) &
                           ** (1.0_D / bet) 
        end function f_pd

        !%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        ! The log of the Gaussian smoothed variance calculated
        ! from my Chebyshev fitting function
        real(D) function logApproxGaussSmoothedVariance(root, logR)

                use mUniverse
                use mRoot

                type(tRoot), intent(inout) :: root
                real(D), intent(in) :: logR

                real(D) :: a
                type(tSmith) :: smith

                a = getParam(root)
                smith = getSmith(root)
                logApproxGaussSmoothedVariance = growthFactor(getUniverse( getSmith(root) ), a )
                logApproxGaussSmoothedVariance = log(logApproxGaussSmoothedVariance ** 2 * evaluate( smith%GSV, exp(logR) ))

        end function logApproxGaussSmoothedVariance

end module mPSFittingFunctions
